/*---------------------------------------------------------------------------*\
|* This file is part of the koala project with MIT License.                  *|
|* Copyright (c) James <zhuguangxiang@gmail.com>                             *|
\*--------------------------------------------------------------------------**/

/**
 * The Koalal IR is register-memory(RegMem) model. This means local
 * varibles(include parameters) of function and global variables are regarded to
 * be stored in memory. And all vars must be loaded into registers and then be
 * computed. It's like ARM CPU load-store architecture and like LLVM IR.
 */

/*

The source of Koala language:

func foo(a int, b int) int {
    var c = 100
    c = a + b
    return c
}

The Koala IR geneated:

func foo(param %a int, param %b int) int {
    local %c int
    set %c, 100 int
    %0 int = get %a
    %1 int = get %b
    %2 int = add %0, %1
    set %c, %2
    %3 int = get %c
    ret %3
}

After optimizations(remove useless set_local) and LSRA, to Koala VM byte codes:

func foo(%0 int, %1 int) int {
    add %0, %0, %1
    ret %0
}
*/

#include "ir.h"
#include "util/log.h"

#ifdef __cplusplus
extern "C" {
#endif

/* types */

static KlrType type_void = { KLR_TYPE_VOID };
static KlrType type_int8 = { KLR_TYPE_INT8 };
static KlrType type_int16 = { KLR_TYPE_INT16 };
static KlrType type_int32 = { KLR_TYPE_INT32 };
static KlrType type_int64 = { KLR_TYPE_INT64 };
static KlrType type_float32 = { KLR_TYPE_FLOAT32 };
static KlrType type_float64 = { KLR_TYPE_FLOAT64 };
static KlrType type_bool = { KLR_TYPE_BOOL };
static KlrType type_char = { KLR_TYPE_CHAR };
static KlrType type_string = { KLR_TYPE_STRING };

KlrType *klr_type_void(void) { return &type_void; }
KlrType *klr_type_int8(void) { return &type_int8; }
KlrType *klr_type_int16(void) { return &type_int16; }
KlrType *klr_type_int32(void) { return &type_int32; }
KlrType *klr_type_int64(void) { return &type_int64; }
KlrType *klr_type_float32(void) { return &type_float32; }
KlrType *klr_type_float64(void) { return &type_float64; }
KlrType *klr_type_bool(void) { return &type_bool; }
KlrType *klr_type_char(void) { return &type_char; }
KlrType *klr_type_string(void) { return &type_string; }
KlrType *klr_type_proto(KlrType *ret, KlrType *params[], int num_params)
{
    KlrProtoType *proto = mm_alloc_obj(proto);
    proto->kind = KLR_TYPE_PROTO;
    proto->ret = ret;
    if (num_params > 0) {
        Vector *vec = vector_create_ptr();
        for (int i = 0; i < num_params; i++) {
            vector_push_back(vec, params + i);
        }
        proto->params = vec;
    }
    return (KlrType *)proto;
}

static int klr_type_equal(KlrType *ty1, KlrType *ty2) { return 1; }

static KlrConst *new_const(KlrType *ty)
{
    KlrConst *val = mm_alloc_obj(val);
    val->kind = KLR_VALUE_CONST;
    init_list(&val->use_list);
    val->type = ty;
    val->name = "";
    return val;
}

KlrValue *klr_const_int8(int8_t v)
{
    KlrType *ty = klr_type_int8();
    KlrConst *val = new_const(ty);
    val->i8val = v;
    return (KlrValue *)val;
}

KlrValue *klr_const_int16(int16_t v)
{
    KlrType *ty = klr_type_int16();
    KlrConst *val = new_const(ty);
    val->i16val = v;
    return (KlrValue *)val;
}

KlrValue *klr_const_int32(int32_t v)
{
    KlrType *ty = klr_type_int32();
    KlrConst *val = new_const(ty);
    val->i32val = v;
    return (KlrValue *)val;
}

KlrValue *klr_const_int64(int64_t v)
{
    KlrType *ty = klr_type_int64();
    KlrConst *val = new_const(ty);
    val->i64val = v;
    return (KlrValue *)val;
}

KlrValue *klr_const_float32(float v)
{
    KlrType *ty = klr_type_float32();
    KlrConst *val = new_const(ty);
    val->f32val = v;
    return (KlrValue *)val;
}

KlrValue *klr_const_float64(double v)
{
    KlrType *ty = klr_type_float64();
    KlrConst *val = new_const(ty);
    val->f64val = v;
    return (KlrValue *)val;
}

KlrValue *klr_const_bool(int v)
{
    KlrType *ty = klr_type_bool();
    KlrConst *val = new_const(ty);
    val->bval = v;
    return (KlrValue *)val;
}

KlrValue *klr_const_char(uint32_t v)
{
    KlrType *ty = klr_type_char();
    KlrConst *val = new_const(ty);
    val->cval = v;
    return (KlrValue *)val;
}

KlrValue *klr_const_string(char *v)
{
    KlrType *ty = klr_type_string();
    KlrConst *val = new_const(ty);
    val->sval = v;
    return (KlrValue *)val;
}

/* basic blocks & edges */

static KlrBasicBlock *new_block(KlrFunc *fn, char *name)
{
    KlrBasicBlock *bb = mm_alloc_obj(bb);
    bb->kind = KLR_VALUE_BLOCK;
    init_list(&bb->use_list);
    bb->name = name ? name : "";

    init_list(&bb->link);
    bb->func = fn;

    init_list(&bb->local_list);
    init_list(&bb->insn_list);
    init_list(&bb->in_edges);
    init_list(&bb->out_edges);
    return bb;
}

KlrBasicBlock *klr_append_block(KlrValue *fn_val, char *label)
{
    KlrFunc *fn = (KlrFunc *)fn_val;
    KlrBasicBlock *bb = new_block(fn, label);

    if (list_empty(&fn->bb_list)) {
        /* first block, add an edge <start, bb> */
        klr_link_age(fn->sbb, bb);
    }

    list_push_back(&fn->bb_list, &bb->link);

    return bb;
}

KlrBasicBlock *klr_add_block(KlrBasicBlock *bb, char *label)
{
    KlrFunc *fn = bb->func;
    KlrBasicBlock *_bb = new_block(fn, label);
    list_add(&bb->link, &_bb->link);
    return _bb;
}

KlrBasicBlock *klr_add_block_before(KlrBasicBlock *bb, char *label)
{
    KlrFunc *fn = bb->func;
    KlrBasicBlock *_bb = new_block(fn, label);
    list_add_before(&bb->link, &_bb->link);
    return _bb;
}

void klr_delete_block(KlrBasicBlock *bb)
{
    list_remove(&bb->link);
    mm_free(bb);
}

void klr_link_age(KlrBasicBlock *src, KlrBasicBlock *dst)
{
    KlrEdge *edge = mm_alloc_obj(edge);
    edge->src = src;
    edge->dst = dst;
    init_list(&edge->link);
    init_list(&edge->in_link);
    init_list(&edge->out_link);

    KlrFunc *fn = src->func;
    list_push_back(&fn->edge_list, &edge->link);
    list_push_back(&src->out_edges, &edge->out_link);
    list_push_back(&dst->in_edges, &edge->in_link);
}

/* instructions */

static KlrInsn *new_insn(KlrInsnKind opcode, int num_opers, char *name)
{
    KlrInsn *insn = mm_alloc(sizeof(*insn) + sizeof(KlrOper) * num_opers);
    insn->kind = KLR_VALUE_INSN;
    init_list(&insn->use_list);
    insn->name = name ? name : "";
    insn->tag = -1;
    insn->opcode = opcode;
    insn->num_opers = num_opers;
    insn->vreg = -1;
    init_list(&insn->bb_link);
    return insn;
}

static void append_insn(KlrBuilder *bldr, KlrInsn *insn)
{
    KlrBasicBlock *bb = bldr->bb;
    list_add(bldr->it, &insn->bb_link);
    insn->bb = bb;
    bldr->it = &insn->bb_link;
    ++bb->num_insns;

    KlrFunc *func = bb->func;
    ++func->num_insns;
}

static KlrUse *new_use(KlrInsn *insn, KlrValue *val)
{
    KlrUse *use = mm_alloc_obj(use);
    use->insn = insn;
    init_list(&use->use_link);
    list_push_back(&val->use_list, &use->use_link);
    use->parent = val;
    return use;
}

static void init_oper(KlrOper *oper, KlrInsn *insn, KlrValue *ref)
{
    switch (ref->kind) {
        case KLR_VALUE_CONST:
            oper->kind = KLR_OPER_CONST;
            oper->konst = (KlrConst *)ref;
            break;
        case KLR_VALUE_GLOBAL:
            oper->kind = KLR_OPER_GLOBAL;
            oper->use = new_use(insn, ref);
            break;
        case KLR_VALUE_FUNC:
            oper->kind = KLR_OPER_FUNC;
            oper->use = new_use(insn, ref);
            break;
        case KLR_VALUE_BLOCK:
            oper->kind = KLR_OPER_BLOCK;
            oper->use = new_use(insn, ref);
            break;
        case KLR_VALUE_PARAM:
        case KLR_VALUE_LOCAL:
            oper->kind = KLR_OPER_LOCAL;
            oper->use = new_use(insn, ref);
            break;
        case KLR_VALUE_INSN:
            oper->kind = KLR_OPER_REG;
            oper->use = new_use(insn, ref);
            break;
        default:
            panic("Invalid kind %d of KlrValue", ref->kind);
            break;
    }
}

KlrValue *klr_get_param(KlrValue *fn_val, int index)
{
    KlrFunc *func = (KlrFunc *)fn_val;

    int size = vector_size(&func->params);
    if (index < 0 || index >= size) {
        panic("index %d out of range(0 ..< %d)\n", index, size);
    }

    KlrValue *param = NULL;
    vector_get(&func->params, index, &param);
    return param;
}

static KlrLocal *new_local(KlrType *ty, char *name)
{
    KlrLocal *local = mm_alloc_obj(local);
    local->kind = KLR_VALUE_LOCAL;
    init_list(&local->use_list);
    local->name = name ? name : "";
    local->tag = -1;
    local->type = ty;
    init_list(&local->bb_link);
    return local;
}

KlrValue *klr_add_local(KlrBuilder *bldr, KlrType *ty, char *name)
{
    KlrLocal *local = new_local(ty, name);
    KlrBasicBlock *bb = bldr->bb;
    list_push_back(&bb->local_list, &local->bb_link);
    local->bb = bb;
    // local->vreg = bb->func->vregs++;
    KlrFunc *func = bb->func;
    vector_push_back(&func->locals, &local);
    return (KlrValue *)local;
}

/**
 * IR: set_local %var, %val
 * %var is local or parameter of function
 * %val is constant or instruction which has result
 * NOTE: %val which is a local var is not allowed. This is needed to use
 * get_local and then set_local like below:
 *
 * var foo = 100
 * var bar = foo
 *
 * local %foo int
 * set %foo, 100 int
 *
 * %0 int = get_local %foo
 * set %bar, %0 int
 */
void klr_build_set(KlrBuilder *bldr, KlrValue *var, KlrValue *val)
{
    if (var->kind != KLR_VALUE_GLOBAL && var->kind != KLR_VALUE_LOCAL) {
        panic("'set %%x, %%v' requires a local/global var.");
    }

    if (val->kind != KLR_VALUE_CONST && val->kind != KLR_VALUE_PARAM &&
        val->kind != KLR_VALUE_INSN) {
        panic("'set_local %%x, %%v' requires a reg value.");
    }

    if (!klr_type_equal(var->type, val->type)) {
        panic("'set_local %%x, %%v' requires the same types.");
    }

    KlrInsn *insn = new_insn(KLR_INSN_SET, 2, "");
    init_oper(&insn->opers[0], insn, var);
    init_oper(&insn->opers[1], insn, val);
    append_insn(bldr, insn);
}

/**
 * IR: %0 int = get_local %foo
 * %foo is a local or parameter of function
 *
 */
KlrValue *klr_build_get(KlrBuilder *bldr, KlrValue *var)
{
    if (var->kind != KLR_VALUE_GLOBAL && var->kind != KLR_VALUE_LOCAL) {
        panic("'get %%x, %%v' requires a local/global var.");
    }

    KlrInsn *insn = new_insn(KLR_INSN_GET, 1, "");
    init_oper(&insn->opers[0], insn, var);
    insn->type = var->type;
    append_insn(bldr, insn);
    return (KlrValue *)insn;
}

static int binop_is_bool(KlrInsnKind op)
{
    static KlrInsnKind ops[] = {
        KLR_INSN_CMPEQ, KLR_INSN_CMPNE, KLR_INSN_CMPGT, KLR_INSN_CMPGE,
        KLR_INSN_CMPLT, KLR_INSN_CMPLE, KLR_INSN_LAND,  KLR_INSN_LOR,
    };

    for (int i = 0; i < COUNT_OF(ops); i++) {
        if (op == ops[i]) return 1;
    }

    return 0;
}

KlrValue *klr_build_binary(KlrBuilder *bldr, KlrValue *lhs, KlrValue *rhs,
                           KlrInsnKind op, char *name)
{
    if (lhs->kind != KLR_VALUE_CONST && lhs->kind != KLR_VALUE_INSN) {
        panic("'add %%x, %%y' requires both reg vars/consts");
    }

    if (rhs->kind != KLR_VALUE_CONST && rhs->kind != KLR_VALUE_INSN) {
        panic("'add %%x, %%y' requires both reg vars/consts");
    }

    KlrInsn *insn = new_insn(op, 2, name);
    init_oper(&insn->opers[0], insn, lhs);
    init_oper(&insn->opers[1], insn, rhs);
    KlrType *ty = lhs->type;
    if (binop_is_bool(op)) ty = klr_type_bool();
    insn->type = ty;
    append_insn(bldr, insn);
    return (KlrValue *)insn;
}

void klr_build_jmp(KlrBuilder *bldr, KlrBasicBlock *dst)
{
    KlrInsn *insn = new_insn(KLR_INSN_JMP, 1, "");
    init_oper(&insn->opers[0], insn, (KlrValue *)dst);
    append_insn(bldr, insn);
    klr_link_age(bldr->bb, dst);
}

void klr_build_condjmp(KlrBuilder *bldr, KlrValue *cond, KlrBasicBlock *_then,
                       KlrBasicBlock *_else)
{
    KlrInsn *insn = new_insn(KLR_INSN_COND_JMP, 3, "");

    init_oper(&insn->opers[0], insn, cond);
    init_oper(&insn->opers[1], insn, (KlrValue *)_then);
    init_oper(&insn->opers[2], insn, (KlrValue *)_else);

    append_insn(bldr, insn);
    klr_link_age(bldr->bb, _then);
    klr_link_age(bldr->bb, _else);
}

void klr_build_ret(KlrBuilder *bldr, KlrValue *ret)
{
    KlrInsn *insn = new_insn(KLR_INSN_RET, 1, "");
    init_oper(&insn->opers[0], insn, ret);
    append_insn(bldr, insn);
    KlrFunc *fn = bldr->bb->func;
    klr_link_age(bldr->bb, fn->ebb);
}

void klr_build_ret_void(KlrBuilder *bldr)
{
    KlrInsn *insn = new_insn(KLR_INSN_RET_VOID, 0, "");
    append_insn(bldr, insn);
    KlrFunc *fn = bldr->bb->func;
    klr_link_age(bldr->bb, fn->ebb);
}

/* modules & functions & classes etc */

KlrModule *klr_create_module(char *name)
{
    KlrModule *m = mm_alloc_obj(m);
    m->name = name;
    vector_init_ptr(&m->variables);
    vector_init_ptr(&m->functions);
    return m;
}

void klr_destroy_module(KlrModule *m) { mm_free(m); }

KlrValue *klr_add_func(KlrModule *m, KlrType *ty, char *name)
{
    KlrFunc *fn = mm_alloc_obj(fn);
    fn->kind = KLR_VALUE_FUNC;
    init_list(&fn->use_list);
    fn->name = name ? name : "";
    fn->type = ty;

    init_list(&fn->bb_list);
    init_list(&fn->edge_list);
    vector_init_ptr(&fn->params);
    vector_init_ptr(&fn->locals);

    /* initial 'start' and 'end' block */
    fn->sbb = new_block(fn, "start");
    fn->ebb = new_block(fn, "end");

    /* add params */
    Vector *pty = ((KlrProtoType *)ty)->params;
    KlrType **item;
    vector_foreach(item, pty, {
        KlrParam *val = mm_alloc_obj(val);
        val->kind = KLR_VALUE_PARAM;
        init_list(&val->use_list);
        val->name = "";
        val->tag = -1;
        val->type = *item;
        vector_push_back(&fn->params, &val);
    });

    vector_push_back(&m->functions, &fn);
    return (KlrValue *)fn;
}

KlrValue *klr_add_global(KlrModule *m, KlrType *ty, char *name) { return NULL; }

#ifdef __cplusplus
}
#endif
